package com.example.module.ai.api;

import cn.hutool.core.util.StrUtil;
import com.example.framework.security.core.util.SecurityFrameworkUtils;
import com.example.module.ai.tencent.websocket.DeepSeekAppTencentTokenService;
import com.example.module.ai.tencent.websocket.TencentAppSession;
import io.socket.client.IO;
import io.socket.client.Socket;
import io.socket.engineio.client.transports.WebSocket;
import jakarta.annotation.Resource;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.json.JSONException;
import org.json.JSONObject;
import org.springframework.stereotype.Service;

import java.net.URISyntaxException;
import java.util.*;

import static java.util.Collections.singletonMap;

@Slf4j
@Service
// socket.io-client-java docs: https://socketio.github.io/socket.io-client-java/#
public class DeepSeekAppTencentWebSocketAiApiImpl implements AiApi {

    private static final String webSocketUrl = "wss://wss.lke.cloud.tencent.com"; // websocket 服务器的 URL
    private static final String path = "/v1/qbot/chat/conn/"; // 服务路径

    @Resource
    private DeepSeekAppTencentTokenService deepSeekAppTencentTokenService;

    private static final Map<String, Map<String, Socket>> USER_SESSION_MAP = new HashMap<>();

    private static final Map<String, Map<String, List<String>>> ANSWER_MAP = new HashMap<>();

    @Override
    public String completion(String question) {
        return completion(question, null, null);
    }

    @SneakyThrows
    @Override
    public String completion(String content, String sid, String reqId) {
        String loginUserId = SecurityFrameworkUtils.getLoginUserId();

        if (StrUtil.isBlank(sid)) {
            sid = TencentAppSession.getSessionId();
            log.info("generated session id: {}, userId: {}", sid, loginUserId);
        }
        if (StrUtil.isBlank(reqId)) {
            reqId = TencentAppSession.getRequestId();
            log.info("generated request id: {}, userId: {}", reqId, loginUserId);
        }

        Socket socket = getSocketByUserIdAndSessionId(loginUserId, sid);

        JSONObject payload = new JSONObject();
        payload.put("request_id", reqId);
        payload.put("content", content);
        payload.put("session_id", sid);
        JSONObject data = new JSONObject();
        data.put("payload", payload);
        // send消息
        socket.emit("send", data);

        // 等待回复
        synchronized (socket) {
            socket.wait();
        }
        return ANSWER_MAP.get(loginUserId).get(sid).getLast();
    }

    @SneakyThrows
    private Socket getSocketByUserIdAndSessionId(String loginUserId, String sid) {
        String token = deepSeekAppTencentTokenService.getWsToken();
        log.info("token:{}", token);

        Socket socket;
        Map<String, Socket> sessionMap = USER_SESSION_MAP.get(loginUserId);
        if (Objects.isNull(sessionMap)) {
            socket = initSocket(token, loginUserId, sid);
            USER_SESSION_MAP.put(loginUserId, new HashMap<>(Map.of(sid, socket)));
        } else {
            socket = sessionMap.get(sid);
            if (Objects.isNull(socket)) {
                socket = initSocket(token, loginUserId, sid);
                sessionMap.put(sid, socket);
            }
        }
        return socket;
    }

    private Socket initSocket(String token, String userId, String sid) throws URISyntaxException {
        IO.Options options = IO.Options.builder()
                .setQuery("EIO=4")
                .setPath(path)  // path需要以'/'开头
                .setTransports(new String[]{WebSocket.NAME})
                .setAuth(singletonMap("token", token))
                .setTimeout(2000)
                .build();

        // 创建 Socket.IO 客户端实例
        Socket socket = IO.socket(webSocketUrl, options);

        // 监听连接成功事件
        socket.on(Socket.EVENT_CONNECT, args1 -> log.info("EVENT_CONNECT: {}", args1))
                .on(Socket.EVENT_CONNECT_ERROR, args2 -> {
                    log.info("EVENT_CONNECT_ERROR: {}", args2);
                    // 如果连接失败，可能需要更新auth，再发起连接
                    options.auth.put("token", deepSeekAppTencentTokenService.getWsToken());
                    socket.disconnect().connect();
                }) // 监听连接错误事件
                .on(Socket.EVENT_DISCONNECT, args3 -> log.info("EVENT_DISCONNECT: {}", args3)) // 监听连接断开事件
                .on("reply", args4 -> {
//                log.info("length:" + args4.length + ", Received reply: " + args4));
                    if (args4.length > 0) {
                        JSONObject reply = (JSONObject) args4[0];
                        try {
                            JSONObject payload = reply.getJSONObject("payload");
                            boolean isFromSelf = payload.getBoolean("is_from_self");
                            boolean isFinal = payload.getBoolean("is_final");
                            String respContent = payload.getString("content");
                            if (isFromSelf) {
                                log.info("Received reply, is_from_self, content: {}", respContent);
                            } else if (isFinal) {
                                log.info("Received reply, is_final, content: {}", respContent);
                                setAnswer(respContent, userId, sid);
                                // 通知已完成回答
                                synchronized (socket) {
                                    socket.notifyAll();
                                }
                            }
                        } catch (JSONException e) {
                            log.error("e:{}", e.getMessage());
                        }
                    } else {
                        log.info("Received reply, but args4.length <= 0");
                    }
                }) // 监听reply事件
                .on("error", args5 -> log.info("Received error: {}", args5)); // 监听error事件
        socket.connect();

        return socket;
    }

    private synchronized void setAnswer(String answer, String userId, String sid) {
        Map<String, List<String>> sessionMap = ANSWER_MAP.get(userId);
        if (Objects.isNull(sessionMap)) {
            ANSWER_MAP.put(userId, new HashMap<>(Map.of(sid, new ArrayList<>(List.of(answer)))));
        } else {
            List<String> answers = sessionMap.get(sid);
            if (Objects.isNull(answers)) {
                sessionMap.put(sid, new ArrayList<>(List.of(answer)));
            } else {
                answers.add(answer);
            }
        }
    }
}